#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include "io.h"
#include "operator.h"

#define FILE_NUM 4

int main(int argc, char** argv)
{
    char* matrix_file_prefix = "A_coo_bin";
    char* rhs_file           = "b_bin";
    char  matrix_file[512];

    // COO matrix
    int      m[FILE_NUM], n[FILE_NUM], num_nonzeros[FILE_NUM], global_rows, global_col, **ii, **jj;
    double** val;

    ii  = (int**)malloc(FILE_NUM * sizeof(int*));
    jj  = (int**)malloc(FILE_NUM * sizeof(int*));
    val = (double**)malloc(FILE_NUM * sizeof(double*));

    // CSR matrix
    int     row, col, nnz, *ia, *ja;
    double* aa;

    // rhs
    int     n_b, i, k;
    double *b, *x;

    // timer
    double tt, time;

    /* ========================================== */
    // Step 1: Load matrix and rhs from binary files
    /* ========================================== */
    printf("Reading COO matrix from multi-binary files %s\n", matrix_file_prefix);

    for (i = 0; i < FILE_NUM; i++) {
        sprintf(matrix_file, "%s.%d", matrix_file_prefix, i);
        printf("reading %s\n", matrix_file);

        // read the matrix Ai from the i-th  file
        COOMatrixRead_Binary(matrix_file, &global_rows, &global_col, &m[i], &n[i], &num_nonzeros[i], &ii[i], &jj[i], &val[i]);
        // printf("%d %d %d\n", m[i], n[i], num_nonzeros[i]);
        // for (k = 0; k < num_nonzeros[i]; k++) {
        //     printf("%d %d %e\n", ii[i][k], jj[i][k], val[i][k]);
        // }
    }

    // sum
    int sum_rows = 0, sum_cols = 0, sum_nonzeros = 0;
    for (i = 0; i < FILE_NUM; i++) {
        sum_rows += m[i];
        sum_cols += n[i];
        sum_nonzeros += num_nonzeros[i];
    }

    // check size
    if (sum_rows != global_rows || sum_cols != global_col) {
        printf("Error: the matrix has problem.\n");
        return 1;
    }

    //! COO to CSR, note that CSR is a global matrix and is only used as an example for small-scale matrices.
    row = sum_rows;
    col = sum_cols;
    nnz = sum_nonzeros;
    printf("row: %d, col: %d, nnz: %d\n", row, col, nnz);

    ia = (int*)malloc((row + 1) * sizeof(int));
    ja = (int*)malloc(nnz * sizeof(int));
    aa = (double*)malloc(nnz * sizeof(double));

    // Initialize row pointer
    for (i = 0; i <= row; i++) {
        ia[i] = 0;
    }

    // Count the number of non-zero entries in each row
    for (i = 0; i < FILE_NUM; i++) {
        for (k = 0; k < num_nonzeros[i]; k++) {
            ia[ii[i][k] + 1]++;
        }
    }

    // Cumulative sum of row pointer
    for (i = 1; i <= row; i++) {
        ia[i] += ia[i - 1];
    }

    // Populate column indices and non-zero values
    for (i = 0; i < FILE_NUM; i++) {
        for (k = 0; k < num_nonzeros[i]; k++) {
            int row_ind   = ii[i][k];
            int destIndex = ia[row_ind]++;
            ja[destIndex] = jj[i][k];
            aa[destIndex] = val[i][k];
        }
    }

    // Restore the row pointer
    for (i = row; i > 0; i--) {
        ia[i] = ia[i - 1];
    }
    ia[0] = 0;

    // free memory
    for (i = 0; i < FILE_NUM; i++) {
        free(ii[i]);
        free(jj[i]);
        free(val[i]);
    }
    free(ii);
    free(jj);
    free(val);

    CSRMatrixPrint("A_check.dat", row, col, nnz, ia, ja, aa);

    printf("Reading right-hand side vector from a binary file %s\n", rhs_file);
    VectorRead_Binary(rhs_file, &b, &n_b);
    // n_b = row;
    // b   = (double*)malloc(sizeof(double) * n_b);
    // for (i = 0; i < n_b; i++) b[i] = 1.0;

    if (n_b != row) {
        printf("Error: The number of rows in the matrix does not match the number of rows in the right-hand side vector.\n");
        return 1;
    }
    x = (double*)malloc(sizeof(double) * n_b);
    for (i = 0; i < n_b; i++) x[i] = 0.0;

    /* ========================================== */
    // Step 2: Solve the linear system
    /* ========================================== */
    tt = GetCurrentTime(); // Start timing
    printf("\nStart solving linear systems ...\n");
    //! Contestants add solution methods here
    // ......
    //
    // Example:
    // CSR_solver(row, col, nnz, ia, ja, aa, b, x);
    // ......

    /* ========================================== */
    // Step 3: Check time and correctness
    /* ========================================== */
    time = (GetCurrentTime() - tt); // End timing
    printf("\nStart checking ...\n");
    fprintf(stdout, "Check total time :              %12.6lf ms\n", time);
    check_correctness(row, ia, ja, aa, x, b);

    // store x to a file
    char* answer_x = "answer_x.dat";
    store_x(n_b, x, answer_x);
    printf("End!\n");

    // free memory
    free(ia);
    free(ja);
    free(aa);
    free(b);
    free(x);

    return 0;
}